# Logistic回归，处理分类问题

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
import numpy as np
import matplotlib.pyplot as plt

# 1、构建数据集

# 加载 torchvision 中包含的数据集，'../dataset/minst'存放在本地的位置，train=True是否是训练集，download=True如果本地没有是否下载
# train_set = torchvision.datasets.MNIST('../dataset/minst', train=True, download=True)
# 加载 MNIST 中的测试集，和上面只有 train=False 不同
# test_set = torchvision.datasets.MNIST('../dataset/minst', train=False, download=True)

# 定义数据集，这里练习使用,第一层中括号里面包含的是单个样本，第二层里面是大哥样本的特征，只有一个特征维度
x_data = torch.Tensor([[1.0], [2.0], [3.0]])
y_data = torch.Tensor([[0], [0], [1]])


# 2、定义模型，继承torch.nn.Module
# 2.1、实现 __init__ 函数，用于变量的初始化，第一行调用父类的 __init__
# 2.2、定义前馈计算 forward() 函数

class LogisticRegressionModel(nn.Module):
    def __init__(self):
        """
        进行初始化
        """
        # 调用父类的 __init__ 方法
        super(LogisticRegressionModel, self).__init__()
        # 定义线性模型，nn.Linear(in_features,out_features,bias) 分别是输入样本的特征维度，输出结果维度，是否使用偏执项，默认使用
        self.linear = nn.Linear(1, 1)

    # 进行前馈计算，forward是一个callable函数，在父类中实现，这里是对父类中 forward 函数的重载
    def forward(self, x):
        # 获得经过 sigmoid 函数之后的输出值
        y_pred = torch.sigmoid(self.linear(x))
        # 过时了，使用上面的方法
        # y_pred = F.sigmoid(self.linear(x))
        return y_pred


# 获得模型结构
model = LogisticRegressionModel()

# 3、定义损失函数和优化器
# 定义损失函数，使用二分类交叉熵损失函数作为损失函数，不去平均值
criterion = torch.nn.BCELoss(size_average=False)
# 定义优化器 optimizer
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

# 4、 训练，定义训练轮数等信息
for epoch in range(1000):
    # 计算预测值 y_pred
    y_pred = model(x_data)
    # 计算损失，loss是一个张量，item()返回其标量值
    loss = criterion(y_pred, y_data)
    print(epoch, loss.item())
    # 梯度清 0
    optimizer.zero_grad()
    # 反向传播求梯度
    loss.backward()
    # 参数更新
    optimizer.step()

# 绘图
x = np.linspace(0, 10, 200)
x_t = torch.Tensor(x).view((200, 1))
# 获得预测值
y_t = model(x_t)
y = y_t.data.numpy()
plt.plot(x, y)
plt.plot([0, 10], [0.5, 0.5], c='r')
plt.xlabel('Hours')
plt.ylabel('Probability of Pass')
plt.grid()
plt.show()
